# -*- coding: utf-8 -*-
"""
------------------------------------------------------------------------------
    File Name:  cifar-10_demo
    Author   :  wanwei1029
    Date     :  2019/1/11
    Desc     :
------------------------------------------------------------------------------
"""
from keras.datasets import cifar10
from keras.utils import np_utils
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.optimizers import SGD, Adam, RMSprop

IMG_CHANNEL = 3
IMG_ROWS = 32
IMG_COLS = 32
BATCH_SIZE = 128
NB_EPOCH = 20
NB_CLASSES = 10
VERBOSE = 2
VALIDATION_SPLIT = 0.2
OPTIMIZER = RMSprop()


def process():
    """
    通过ImageDataGenerator做数据增强，还未处理，可参见：
    https://blog.csdn.net/jacke121/article/details/79245732
    :return:
    """
    (x_train, y_train), (x_test, y_test) = load_cifar10_data()
    y_train = np_utils.to_categorical(y_train, NB_CLASSES)
    y_test = np_utils.to_categorical(y_test, NB_CLASSES)
    x_train = x_train.astype('float32')
    x_test = x_test.astype('float32')
    x_train /= 255
    x_test /= 255
    model = get_model_02()
    model.compile(loss='categorical_crossentropy', optimizer=OPTIMIZER, metrics=['accuracy'])

    history = model.fit(x_train, y_train, batch_size=BATCH_SIZE, epochs=NB_EPOCH, verbose=VERBOSE,
                        validation_split=VALIDATION_SPLIT)

    score = model.evaluate(x_test, y_test, verbose=VERBOSE)

    print("Test score:", score[0])
    print("Test accuracy:", score[1])


def get_model_01():
    """
    简单模型，最终结果，loss:1.0111，accuracy: 0.6754
    :return:
    """
    model = Sequential()
    model.add(Conv2D(32, (3, 3), padding="same", input_shape=(IMG_ROWS, IMG_COLS, IMG_CHANNEL)))
    model.add(Activation("relu"))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))
    model.add(Flatten())
    model.add(Dense(512))
    model.add(Activation("relu"))
    model.add(Dropout(0.5))
    model.add(Dense(NB_CLASSES))
    model.add(Activation("softmax"))
    model.summary()
    return model


def get_model_02():
    """
    模型最终结果：loss:0.7416 ,accuracy: 0.7705
    :return:
    """
    model = Sequential()
    model.add(Conv2D(32, (3, 3), padding="same", input_shape=(IMG_ROWS, IMG_COLS, IMG_CHANNEL)))
    model.add(Activation("relu"))
    model.add(Conv2D(32, (3, 3), padding="same"))
    model.add(Activation("relu"))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))
    model.add(Conv2D(64, (3, 3), padding="same"))
    model.add(Activation("relu"))
    model.add(Conv2D(64, 3, 3))
    model.add(Activation("relu"))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))
    model.add(Flatten())
    model.add(Dense(512))
    model.add(Activation("relu"))
    model.add(Dropout(0.5))
    model.add(Dense(NB_CLASSES))
    model.add(Activation("softmax"))
    model.summary()
    return model


def load_cifar10_data():
    (x_train, y_train), (x_test, y_test) = cifar10.load_data()
    print(x_train.shape)
    print(y_train.shape)
    return (x_train, y_train), (x_test, y_test)


def demo():
    """
    """
    process()


if __name__ == '__main__':
    test_method = "demo"
    if test_method == "demo":
        demo()
